"""
ComfyUI服务
负责与ComfyUI服务器通信，提交工作流和监控任务进度
支持多服务器负载均衡，当有多个任务时选择空闲的服务器进行处理
"""

import json
import asyncio
import logging
from typing import Dict, Any, Optional, AsyncGenerator, List
from datetime import datetime, timedelta
from dataclasses import dataclass

from ..core.http_client import get_http_client
from ..core.config import settings
from .retry_service import async_retry
from .workflow_manager import workflow_manager

logger = logging.getLogger(__name__)


@dataclass
class ComfyUIServer:
    """ComfyUI服务器信息"""
    url: str
    is_healthy: bool = True
    current_tasks: int = 0
    last_check: Optional[datetime] = None
    response_time: float = 0.0


class ComfyUILoadBalancer:
    """ComfyUI负载均衡器"""
    
    def __init__(self, servers: List[str], api_key: str):
        self.api_key = api_key
        self.servers = [ComfyUIServer(url=url) for url in servers]
        self.task_queue = asyncio.Queue()
        self.server_tasks = {}  # 跟踪每个服务器的任务
        
    async def get_server_status(self, server: ComfyUIServer) -> Dict[str, Any]:
        """获取服务器状态"""
        try:
            client = get_http_client()
            start_time = datetime.now()
            
            response = await client.get(
                f"{server.url}/queue",
                headers={"Authorization": f"Bearer {self.api_key}"},
                timeout=5.0
            )
            
            end_time = datetime.now()
            server.response_time = (end_time - start_time).total_seconds()
            server.last_check = end_time
            
            if response.status_code == 200:
                queue_data = response.json()
                running_count = len(queue_data.get("queue_running", []))
                pending_count = len(queue_data.get("queue_pending", []))
                server.current_tasks = running_count + pending_count
                server.is_healthy = True
                return {
                    "healthy": True,
                    "running_tasks": running_count,
                    "pending_tasks": pending_count,
                    "total_tasks": server.current_tasks
                }
            else:
                server.is_healthy = False
                return {"healthy": False, "error": f"HTTP {response.status_code}"}
                
        except Exception as e:
            server.is_healthy = False
            logger.error(f"检查服务器状态失败 {server.url}: {e}")
            return {"healthy": False, "error": str(e)}
    
    async def select_best_server(self) -> Optional[ComfyUIServer]:
        """选择最适合的服务器"""
        # 检查所有服务器状态
        await asyncio.gather(*[
            self.get_server_status(server) for server in self.servers
        ])
        
        # 过滤健康的服务器
        healthy_servers = [s for s in self.servers if s.is_healthy]
        
        if not healthy_servers:
            logger.error("没有可用的ComfyUI服务器")
            return None
        
        # 选择任务最少的服务器
        best_server = min(healthy_servers, key=lambda s: s.current_tasks)
        logger.info(f"选择服务器: {best_server.url} (当前任务数: {best_server.current_tasks})")
        return best_server
    
    async def submit_to_server(self, server: ComfyUIServer, workflow: Dict[str, Any]) -> Optional[str]:
        """向指定服务器提交工作流"""
        try:
            client = get_http_client()
            payload = {"prompt": workflow}
            
            response = await client.post(
                f"{server.url}/prompt",
                headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
                json=payload,
                timeout=30.0
            )
            
            if response.status_code == 200:
                response_data = response.json()
                prompt_id = response_data.get("prompt_id")
                if prompt_id:
                    # 更新服务器任务计数
                    server.current_tasks += 1
                    self.server_tasks[prompt_id] = server.url
                    logger.info(f"任务 {prompt_id} 已提交到服务器 {server.url}")
                return prompt_id
            else:
                logger.error(f"提交失败 {server.url}: {response.status_code} - {response.text}")
                return None
                
        except Exception as e:
            logger.error(f"提交到服务器异常 {server.url}: {e}")
            return None


class ComfyUIService:
    """
    ComfyUI服务
    负责与ComfyUI服务器的所有交互，支持多服务器负载均衡
    """

    def __init__(self):
        # 初始化负载均衡器
        self.load_balancer = ComfyUILoadBalancer(
            servers=settings.comfyui_servers,
            api_key=settings.comfyui_api_key
        )
        
        # 任务排队机制
        self.pending_tasks = None  # 延迟初始化
        self.max_concurrent_tasks = 20  # 总的最大并发任务数
        self.current_tasks = 0
        
        # 后台任务处理器状态
        self._task_processor_running = False
        self._task_processor_started = False
        
        # 基础工作流模板 - 保持原有模板
        self.workflow_template = {
            "3": {
                "inputs": {
                    "seed": 733488244826594,
                    "steps": 20,
                    "cfg": 8,
                    "sampler_name": "euler",
                    "scheduler": "normal",
                    "denoise": 1,
                    "model": ["4", 0],
                    "positive": ["6", 0],
                    "negative": ["7", 0],
                    "latent_image": ["5", 0],
                },
                "class_type": "KSampler",
                "_meta": {"title": "K采样器"},
            },
            "4": {
                "inputs": {"ckpt_name": "25D_v2.safetensors"},
                "class_type": "CheckpointLoaderSimple",
                "_meta": {"title": "Checkpoint加载器（简易）"},
            },
            "5": {
                "inputs": {"width": 512, "height": 512, "batch_size": 1},
                "class_type": "EmptyLatentImage",
                "_meta": {"title": "空Latent图像"},
            },
            "6": {
                "inputs": {
                    "text": "beautiful scenery nature glass bottle landscape, purple galaxy bottle",
                    "speak_and_recognation": {"__value__": [False, True]},
                    "clip": ["4", 1],
                },
                "class_type": "CLIPTextEncode",
                "_meta": {"title": "CLIP文本编码"},
            },
            "7": {
                "inputs": {
                    "text": "text, watermark",
                    "speak_and_recognation": {"__value__": [False, True]},
                    "clip": ["4", 1],
                },
                "class_type": "CLIPTextEncode",
                "_meta": {"title": "CLIP文本编码"},
            },
            "8": {
                "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
                "class_type": "VAEDecode",
                "_meta": {"title": "VAE解码"},
            },
            "9": {
                "inputs": {"filename_prefix": "ComfyUI", "images": ["8", 0]},
                "class_type": "SaveImage",
                "_meta": {"title": "保存图像"},
            },
        }

    async def _ensure_initialized(self):
        """确保异步组件已初始化"""
        if self.pending_tasks is None:
            self.pending_tasks = asyncio.Queue()
        
        if not self._task_processor_started:
            self._task_processor_running = True
            self._task_processor_started = True
            asyncio.create_task(self._background_task_processor())

    def _get_auth_headers(self) -> Dict[str, str]:
        """
        获取带有认证信息的请求头

        Returns:
            Dict[str, str]: 请求头字典
        """
        return {
            "Authorization": f"Bearer {self.load_balancer.api_key}",
            "Content-Type": "application/json",
        }

    def _create_workflow(
        self,
        character_trigger: Optional[str],
        refined_prompt: str,
        style_name: str = "anime_style_1",
    ) -> Dict[str, Any]:
        """
        创建工作流
        使用工作流管理器创建动态工作流

        Args:
            character_trigger: 角色触发词
            refined_prompt: 提炼后的用户提示词
            style_name: 工作流风格名称

        Returns:
            Dict: 完整的工作流配置
        """
        # 使用工作流管理器创建动态工作流
        workflow = workflow_manager.create_workflow(
            style_name=style_name,
            character_trigger=character_trigger,
            user_prompt=refined_prompt,
        )

        if workflow:
            # 验证工作流
            if workflow_manager.validate_workflow(workflow):

                return workflow
            else:
                logger.error("工作流验证失败，使用备用模板")

        # 如果工作流管理器失败，使用备用的硬编码模板
        return self._create_fallback_workflow(character_trigger, refined_prompt)

    def _create_fallback_workflow(
        self, character_trigger: Optional[str], refined_prompt: str
    ) -> Dict[str, Any]:
        """
        创建备用工作流（保持原有逻辑）

        Args:
            character_trigger: 角色触发词
            refined_prompt: 提炼后的用户提示词

        Returns:
            Dict: 完整的工作流配置
        """
        # 深拷贝模板避免修改原模板
        workflow = json.loads(json.dumps(self.workflow_template))

        # 组合用户输入的提示词和角色触发词
        if character_trigger:
            combined_prompt = f"{refined_prompt}, {character_trigger}".strip(", ")
        else:
            combined_prompt = refined_prompt.strip()

        # 设置正面提示词
        workflow["6"]["inputs"]["text"] = combined_prompt

        # 生成随机种子确保每次生成不同的图像
        import random

        workflow["3"]["inputs"]["seed"] = random.randint(1, 2**32 - 1)

        # 移除向ComfyUI模型传递数据的详细日志

        return workflow

    @async_retry(retries=3, delay=2.0, backoff=1.5)
    async def submit_workflow(
        self,
        character_trigger: Optional[str],
        refined_prompt: str,
        style_name: str = "anime_style_1",
    ) -> Optional[str]:
        """
        提交工作流到ComfyUI服务器（支持负载均衡）

        Args:
            character_trigger: 角色触发词
            refined_prompt: 提炼后的用户提示词
            style_name: 工作流风格名称

        Returns:
            Optional[str]: ComfyUI返回的prompt_id，失败时返回None
        """
        try:
            # 确保异步组件已初始化
            await self._ensure_initialized()
            
            # 创建工作流
            workflow = self._create_workflow(
                character_trigger, refined_prompt, style_name
            )

            # 选择最佳服务器
            best_server = await self.load_balancer.select_best_server()
            if not best_server:
                logger.error("没有可用的ComfyUI服务器")
                return None

            # 检查是否需要排队等待
            total_server_tasks = sum(s.current_tasks for s in self.load_balancer.servers if s.is_healthy)
            if total_server_tasks >= self.max_concurrent_tasks:
                logger.info(f"所有服务器任务数({total_server_tasks})已达上限，任务进入排队")
                await self.pending_tasks.put((character_trigger, refined_prompt, style_name))
                return "queued"  # 返回特殊状态表示已排队

            # 提交到选中的服务器
            prompt_id = await self.load_balancer.submit_to_server(best_server, workflow)
            
            if prompt_id:
                self.current_tasks += 1
                logger.info(f"任务 {prompt_id} 已提交，当前总任务数: {self.current_tasks}")
                return prompt_id
            else:
                # 如果第一个服务器失败，尝试其他服务器
                for server in self.load_balancer.servers:
                    if server != best_server and server.is_healthy:
                        prompt_id = await self.load_balancer.submit_to_server(server, workflow)
                        if prompt_id:
                            self.current_tasks += 1
                            logger.info(f"任务 {prompt_id} 已提交到备用服务器 {server.url}")
                            return prompt_id
                
                logger.error("所有服务器都无法处理任务")
                return None

        except Exception as e:
            logger.error(f"提交工作流时发生异常: {e}")
            return None

    async def get_task_status(self, prompt_id: str) -> Dict[str, Any]:
        """
        获取任务状态（支持多服务器）

        Args:
            prompt_id: ComfyUI任务ID

        Returns:
            Dict: 任务状态信息
        """
        try:
            client = get_http_client()
            
            # 查找任务所在的服务器
            task_server_url = self.load_balancer.server_tasks.get(prompt_id)
            
            # 如果知道任务在哪个服务器，直接查询该服务器
            if task_server_url:
                servers_to_check = [s for s in self.load_balancer.servers if s.url == task_server_url]
            else:
                # 如果不知道，查询所有健康的服务器
                servers_to_check = [s for s in self.load_balancer.servers if s.is_healthy]
            
            for server in servers_to_check:
                try:
                    # 先检查历史记录（判断是否已完成）
                    history_response = await client.get(
                        f"{server.url}/history/{prompt_id}",
                        headers=self._get_auth_headers(),
                        timeout=10.0,
                    )

                    if history_response.status_code == 200:
                        history_data = history_response.json().get(prompt_id)
                        if history_data:
                            # 任务已完成，减少当前任务计数
                            if prompt_id in self.load_balancer.server_tasks:
                                self.current_tasks = max(0, self.current_tasks - 1)
                                del self.load_balancer.server_tasks[prompt_id]
                            
                            # 查找生成的图片
                            outputs = history_data.get("outputs", {})
                            for node_id, output_data in outputs.items():
                                if "images" in output_data and output_data["images"]:
                                    image_info = output_data["images"][0]
                                    return {
                                        "status": "completed",
                                        "progress": 100,
                                        "image_filename": image_info["filename"],
                                        "image_url": f"{server.url}/view?filename={image_info['filename']}",
                                        "server_url": server.url
                                    }

                    # 检查队列状态
                    queue_response = await client.get(
                        f"{server.url}/queue", 
                        headers=self._get_auth_headers(), 
                        timeout=10.0
                    )

                    if queue_response.status_code == 200:
                        queue_data = queue_response.json()

                        # 检查是否在运行队列中
                        running_tasks = [item[1] for item in queue_data.get("queue_running", [])]
                        if prompt_id in running_tasks:
                            return {
                                "status": "processing", 
                                "progress": 60,
                                "server_url": server.url
                            }

                        # 检查是否在等待队列中
                        pending_tasks = queue_data.get("queue_pending", [])
                        for i, task in enumerate(pending_tasks):
                            if len(task) > 1 and task[1] == prompt_id:
                                # 根据队列位置计算进度
                                progress = max(10, 30 - i * 5)
                                return {
                                    "status": "queued",
                                    "progress": progress,
                                    "queue_position": i + 1,
                                    "server_url": server.url
                                }
                except Exception as e:
                    logger.warning(f"查询服务器 {server.url} 状态失败: {e}")
                    continue

            # 如果在所有服务器都没找到，可能任务已失败或不存在
            return {"status": "failed", "progress": 0, "error": "任务未找到"}

        except Exception as e:
            logger.error(f"获取任务状态失败: {e}")
            return {"status": "error", "progress": 0, "error": str(e)}

    async def monitor_task(
        self, prompt_id: str, max_wait_time: int = 600, check_interval: float = 5.0
    ) -> AsyncGenerator[Dict[str, Any], None]:
        """
        监控任务进度
        异步生成器，持续返回任务状态直到完成或超时

        Args:
            prompt_id: ComfyUI任务ID
            max_wait_time: 最大等待时间（秒）
            check_interval: 检查间隔（秒）

        Yields:
            Dict: 任务状态更新
        """
        start_time = asyncio.get_event_loop().time()

        while True:
            elapsed_time = asyncio.get_event_loop().time() - start_time

            # 检查超时
            if elapsed_time > max_wait_time:
                yield {
                    "status": "timeout",
                    "progress": 0,
                    "error": f"任务超时，已等待{max_wait_time}秒",
                }
                break

            # 获取当前状态
            status_info = await self.get_task_status(prompt_id)
            yield status_info

            # 如果任务完成或失败，停止监控
            if status_info.get("status") in ["completed", "failed", "error", "timeout"]:
                break

            # 等待下次检查
            await asyncio.sleep(check_interval)

    async def cancel_task(self, prompt_id: str) -> bool:
        """
        取消任务（支持多服务器）

        Args:
            prompt_id: 要取消的任务ID

        Returns:
            bool: 是否成功取消
        """
        try:
            client = get_http_client()
            
            # 查找任务所在的服务器
            task_server_url = self.load_balancer.server_tasks.get(prompt_id)
            
            if task_server_url:
                # 在指定服务器上取消任务
                servers_to_try = [s for s in self.load_balancer.servers if s.url == task_server_url]
            else:
                # 在所有服务器上尝试取消
                servers_to_try = self.load_balancer.servers
            
            for server in servers_to_try:
                try:
                    response = await client.delete(
                        f"{server.url}/queue",
                        headers=self._get_auth_headers(),
                        json={"delete": [prompt_id]},
                        timeout=10.0,
                    )

                    if response.status_code == 200:
                        # 清理任务记录
                        if prompt_id in self.load_balancer.server_tasks:
                            del self.load_balancer.server_tasks[prompt_id]
                        self.current_tasks = max(0, self.current_tasks - 1)
                        logger.info(f"任务 {prompt_id} 已在服务器 {server.url} 上取消")
                        return True
                except Exception as e:
                    logger.warning(f"在服务器 {server.url} 取消任务失败: {e}")
                    continue

            return False

        except Exception as e:
            logger.error(f"取消任务时发生异常: {e}")
            return False

    async def get_image_url(self, filename: str, server_url: str = None) -> str:
        """
        获取图像的可访问URL

        Args:
            filename: 图像文件名
            server_url: 指定的服务器URL，如果为None则使用第一个健康的服务器

        Returns:
            str: 图像的完整URL
        """
        if server_url:
            return f"{server_url}/view?filename={filename}"
        
        # 如果没有指定服务器，使用第一个健康的服务器
        healthy_servers = [s for s in self.load_balancer.servers if s.is_healthy]
        if healthy_servers:
            return f"{healthy_servers[0].url}/view?filename={filename}"
        
        # 如果没有健康的服务器，使用第一个服务器
        if self.load_balancer.servers:
            return f"{self.load_balancer.servers[0].url}/view?filename={filename}"
        
        raise ValueError("没有可用的ComfyUI服务器")

    async def check_server_health(self) -> Dict[str, Any]:
        """
        检查所有ComfyUI服务器健康状态

        Returns:
            Dict: 所有服务器的健康状态信息
        """
        health_status = {}
        
        for server in self.load_balancer.servers:
            try:
                status = await self.load_balancer.get_server_status(server)
                health_status[server.url] = {
                    "healthy": status["healthy"],
                    "response_time": server.response_time,
                    "current_tasks": server.current_tasks,
                    "last_check": server.last_check.isoformat() if server.last_check else None
                }
                if not status["healthy"]:
                    health_status[server.url]["error"] = status.get("error", "Unknown error")
            except Exception as e:
                health_status[server.url] = {
                    "healthy": False,
                    "error": str(e),
                    "response_time": 0,
                    "current_tasks": 0
                }
        
        return health_status

    async def get_load_balancer_stats(self) -> Dict[str, Any]:
        """
        获取负载均衡器统计信息

        Returns:
            Dict: 负载均衡器状态信息
        """
        # 确保异步组件已初始化
        await self._ensure_initialized()
        
        total_tasks = sum(server.current_tasks for server in self.load_balancer.servers)
        healthy_servers = [s for s in self.load_balancer.servers if s.is_healthy]
        
        return {
            "total_servers": len(self.load_balancer.servers),
            "healthy_servers": len(healthy_servers),
            "total_tasks_across_servers": total_tasks,
            "current_app_tasks": self.current_tasks,
            "pending_queue_size": self.pending_tasks.qsize() if self.pending_tasks else 0,
            "max_concurrent_tasks": self.max_concurrent_tasks,
            "servers": [
                {
                    "url": server.url,
                    "healthy": server.is_healthy,
                    "current_tasks": server.current_tasks,
                    "response_time": server.response_time
                }
                for server in self.load_balancer.servers
            ]
        }

    async def process_pending_tasks(self):
        """
        处理排队中的任务
        """
        # 确保异步组件已初始化
        await self._ensure_initialized()
        
        processed_count = 0
        max_process_per_cycle = 5  # 每次最多处理5个任务，避免无限循环
        
        while (not self.pending_tasks.empty() and 
               processed_count < max_process_per_cycle):
            try:
                # 检查是否有空闲服务器
                best_server = await self.load_balancer.select_best_server()
                if not best_server:
                    break
                
                total_server_tasks = sum(s.current_tasks for s in self.load_balancer.servers if s.is_healthy)
                if total_server_tasks >= self.max_concurrent_tasks:
                    break
                
                character_trigger, refined_prompt, style_name = await asyncio.wait_for(
                    self.pending_tasks.get(), timeout=1.0
                )
                
                # 直接提交到服务器，避免递归调用submit_workflow
                workflow = self._create_workflow(character_trigger, refined_prompt, style_name)
                prompt_id = await self.load_balancer.submit_to_server(best_server, workflow)
                
                if prompt_id:
                    self.current_tasks += 1
                    logger.info(f"排队任务 {prompt_id} 已成功提交到服务器 {best_server.url}")
                    processed_count += 1
                else:
                    logger.warning("排队任务提交失败，重新加入队列")
                    await self.pending_tasks.put((character_trigger, refined_prompt, style_name))
                    break
                    
            except asyncio.TimeoutError:
                break
            except Exception as e:
                logger.error(f"处理排队任务时出错: {e}")
                break

    async def _background_task_processor(self):
        """
        后台任务处理器，定期检查并处理排队的任务
        """
        logger.info("后台任务处理器已启动")
        while self._task_processor_running:
            try:
                await asyncio.sleep(5)  # 每5秒检查一次
                await self.process_pending_tasks()
            except Exception as e:
                logger.error(f"后台任务处理器异常: {e}")
                await asyncio.sleep(10)  # 出错时等待更长时间
        logger.info("后台任务处理器已停止")

    def stop_task_processor(self):
        """
        停止后台任务处理器
        """
        self._task_processor_running = False

    def get_workflow_info(
        self,
        character_trigger: Optional[str],
        refined_prompt: str,
        style_name: str = "anime_style_1",
    ) -> Dict[str, Any]:
        """
        获取工作流信息（用于调试）

        Args:
            character_trigger: 角色触发词
            refined_prompt: 提炼后的用户提示词
            style_name: 工作流风格名称

        Returns:
            Dict: 工作流配置信息
        """
        # 优先使用工作流管理器获取信息
        workflow_info = workflow_manager.get_workflow_info(style_name)
        if workflow_info:
            workflow_info.update(
                {
                    "character_trigger": character_trigger,
                    "refined_prompt": refined_prompt,
                    "available_servers": [s.url for s in self.load_balancer.servers],
                    "healthy_servers": [s.url for s in self.load_balancer.servers if s.is_healthy],
                    "style_name": style_name,
                }
            )
            return workflow_info

        # 备用：创建工作流并提取信息
        workflow = self._create_workflow(character_trigger, refined_prompt, style_name)
        return {
            "workflow_nodes": len(workflow),
            "character_trigger": character_trigger,
            "refined_prompt": refined_prompt,
            "combined_prompt": (
                f"{refined_prompt}, {character_trigger}".strip(", ")
                if character_trigger
                else refined_prompt.strip()
            ),
            "image_size": f"{workflow['5']['inputs']['width']}x{workflow['5']['inputs']['height']}",
            "sampling_steps": workflow["3"]["inputs"]["steps"],
            "seed": workflow["3"]["inputs"]["seed"],
            "available_servers": [s.url for s in self.load_balancer.servers],
            "healthy_servers": [s.url for s in self.load_balancer.servers if s.is_healthy],
            "style_name": style_name,
        }

    def get_available_styles(self) -> List[str]:
        """获取所有可用的工作流风格"""
        return workflow_manager.get_available_styles()


# 创建全局服务实例
comfyui_service = ComfyUIService()
